/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 and
 * only version 2 as published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * Description:
 * Author: huawei
 * Create: 2022-11-29
 */

#include <linux/sched.h>

#include "svm_log.h"
#include "svm_rbtree.h"

int devmm_rb_erase(struct rb_root *root, struct rb_node *node)
{
    int ret = -ENODEV;

    if (RB_EMPTY_NODE(node) == false) {
        rb_erase(node, root);
        RB_CLEAR_NODE(node);
        ret = 0;
    }
    return ret;
}

int devmm_rb_insert(struct rb_root *root, struct rb_node *node, rb_handle_func get_handle)
{
    struct rb_node **cur_node = &root->rb_node;
    struct rb_node *parent = NULL;
    u64 handle = get_handle(node);

    /* Figure out where to put new node */
    while (*cur_node) {
        u64 tmp_handle = get_handle(*cur_node);

        parent = *cur_node;
        if (handle < tmp_handle) {
            cur_node = &((*cur_node)->rb_left);
        } else if (handle > tmp_handle) {
            cur_node = &((*cur_node)->rb_right);
        } else {
            devmm_drv_err("Insert same priv.\n");
            return -EINVAL;
        }
    }

    /* Add new node and rebalance tree. */
    rb_link_node(node, parent, cur_node);
    rb_insert_color(node, root);
    return 0;
}

struct rb_node *devmm_rb_search(struct rb_root *root, u64 handle, rb_handle_func get_handle)
{
    struct rb_node *node = NULL;

    node = root->rb_node;
    while (node != NULL) {
        u64 tmp_handle = get_handle(node);
        if (handle < tmp_handle) {
            node = node->rb_left;
        } else if (handle > tmp_handle) {
            node = node->rb_right;
        } else {
            return node;
        }
    }

    return NULL;
}

/* condition can be NULL */
struct rb_node *devmm_rb_erase_one_node(struct rb_root *root, rb_erase_condition condition)
{
    struct rb_node *node = NULL;

    if (RB_EMPTY_ROOT(root) == true) {
        return NULL;
    }

    node = rb_first(root);
    while (node != NULL) {
        if (condition == NULL) {
            break;
        }

        if (condition(node) == true) {
            break;
        }
        node = rb_next(node);
    }

    if (node != NULL) {
        rb_erase(node, root);
        RB_CLEAR_NODE(node);
    }

    return node;
}
